Refactor Wan Model Training & Add Wan-VACE Training Support#352
Open
Refactor Wan Model Training & Add Wan-VACE Training Support#352
Conversation
- Ensure `adam_weight_decay` is a float. - Add `tensorboard_dir` parameter for logging. Co-authored-by: martinarroyo <martinarroyo@google.com>
- Conditionally apply dropout only when rate > 0. - Use standard list initialization. - Add rngs parameter to layer_forward (essential for gradient checkpointing with dropout > 0) Co-authored-by: martinarroyo <martinarroyo@google.com>
Replaces the hardcoded learning rate in the optimizer creation with the value from `config.learning_rate`. Co-authored-by: martinarroyo <martinarroyo@google.com>
Co-authored-by: martinarroyo <martinarroyo@google.com>
Co-authored-by: martinarroyo <martinarroyo@google.com>
…lices and different topologies Co-authored-by: martinarroyo <martinarroyo@google.com>
…ling and introduce disable_training_weights, add max_grad_norm and max_abs_grad logging. - Switched timestamp sampling from discrete to continuous. - Add max_grad_norm and max_abs_grad calculation and logging. - Introduced `config.disable_training_weights` to optionally disable mid-point loss weighting. Co-authored-by: martinarroyo <martinarroyo@google.com>
The following key functionalities have been moved from WanTrainer to the new `BaseWanTrainer` ABC: - Initialization and config handling - Scheduler creation - TFLOPs calculation - Core training and evaluation loops (`start_training`, `training_loop`, `eval`) - Abstract methods for checkpointer, data loading, sharding, and step functions. Co-authored-by: martinarroyo <martinarroyo@google.com>
Introduces training support for WAN-VACE models. New files: - train_wan_vace.py: Main training script. - wan_vace_trainer.py: Trainer class for WAN-VACE. - wan_vace_checkpointing_2_1.py: Checkpointing logic for WAN-VACE. Co-authored-by: martinarroyo <martinarroyo@google.com>
entrpn
reviewed
Mar 12, 2026
| # Output directory | ||
| # Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/" | ||
| base_output_directory: "" | ||
| tensorboard_dir: "" |
Collaborator
There was a problem hiding this comment.
tensorboard_dir is created automatically inside the pyconfig. Is there a reason it needs to be in the config?
Collaborator
Author
There was a problem hiding this comment.
You're right, it's indeed no longer needed. It was necessary in an older version of the code, but it's redundant now. I've removed it in the latest commit
Collaborator
|
As this is a fairly large refactor:
|
prishajain1
reviewed
Mar 13, 2026
|
|
||
| params_restore = ocp.args.PyTreeRestore( | ||
| restore_args=jax.tree.map( | ||
| lambda _: ocp.RestoreArgs(restore_type=np.ndarray), |
Collaborator
There was a problem hiding this comment.
Passing restore_type = np.ndarray causes the JAX sharding applied above to be redundant. (JAX sharding cannot work on np.ndarrays). Suggest to make it jax.Array to ensure checkpoint is loaded on host in sharded manner if that's intended
prishajain1
reviewed
Mar 13, 2026
| abstract_tree_structure_params = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, transformer_metadata) | ||
| state = metadatas.wan_state | ||
|
|
||
| def add_sharding_to_struct(leaf_struct, sharding): |
Collaborator
There was a problem hiding this comment.
A safer way to do this to prevent unexpected crashes (for any elements not having shape/dtype attributes):
def add_sharding_to_struct(leaf_struct, sharding):
struct = ocp.utils.to_shape_dtype_struct(leaf_struct)
if hasattr(struct, "shape") and hasattr(struct, "dtype"):
return jax.ShapeDtypeStruct(
shape=struct.shape, dtype=struct.dtype, sharding=sharding
)
return struct
- Address review comments: removed tensorboard_dir because it is created automatically inside the pyconfig. Co-authored-by: martinarroyo <martinarroyo@google.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR introduces several improvements and fixes to the Wan model training, as well as adds support for training Wan-VACE models.
Key changes include:
Bug fixes:
Config updates:
Wan-VACE Support:
New Features: